#!/usr/bin/env python3
"""
网络安全漏洞知识图谱项目启动脚本
负责系统初始化、健康检查、数据预处理
"""

import asyncio
import logging
import os
import sys
import time
import subprocess
from typing import Dict, List
import json
import argparse

import docker
import psutil
import requests
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.table import Table
from rich.panel import Panel


class ProjectManager:
    """项目管理器"""
    
    def __init__(self):
        self.console = Console()
        self.docker_client = docker.from_env()
        self.logger = self._setup_logging()
        
        # 服务配置
        self.services = {
            'redis': {'port': 6379, 'healthcheck': self._check_redis},
            'mongodb': {'port': 27017, 'healthcheck': self._check_mongodb},
            'neo4j': {'port': 7474, 'healthcheck': self._check_neo4j},
            'elasticsearch': {'port': 9200, 'healthcheck': self._check_elasticsearch},
            'kafka': {'port': 9092, 'healthcheck': self._check_kafka},
            'minio': {'port': 9000, 'healthcheck': self._check_minio},
        }
        
        # 应用服务
        self.app_services = {
            'api-server': {'port': 8000, 'healthcheck': self._check_api},
            'crawler-service': {'healthcheck': self._check_crawler},
            'ai-processor': {'healthcheck': self._check_ai_processor},
            'celery-worker': {'healthcheck': self._check_celery_worker},
        }
        
    def _setup_logging(self) -> logging.Logger:
        """设置日志"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('logs/project_manager.log'),
                logging.StreamHandler()
            ]
        )
        return logging.getLogger(__name__)
    
    def check_prerequisites(self) -> bool:
        """检查运行环境"""
        self.console.print("\n[bold blue]🔍 检查运行环境...[/bold blue]")
        
        checks = []
        
        # 检查Docker
        try:
            docker_version = self.docker_client.version()
            checks.append(("Docker", "✅", f"版本 {docker_version['Version']}"))
        except Exception as e:
            checks.append(("Docker", "❌", f"未安装或未启动: {e}"))
            return False
        
        # 检查Docker Compose
        try:
            result = subprocess.run(['docker-compose', '--version'], 
                                 capture_output=True, text=True)
            if result.returncode == 0:
                checks.append(("Docker Compose", "✅", result.stdout.strip()))
            else:
                checks.append(("Docker Compose", "❌", "未安装"))
                return False
        except FileNotFoundError:
            checks.append(("Docker Compose", "❌", "未安装"))
            return False
        
        # 检查系统资源
        memory = psutil.virtual_memory()
        disk = psutil.disk_usage('/')
        cpu_count = psutil.cpu_count()
        
        # 内存检查 (至少8GB)
        if memory.total >= 8 * 1024**3:
            checks.append(("内存", "✅", f"{memory.total // 1024**3}GB 可用"))
        else:
            checks.append(("内存", "⚠️", f"{memory.total // 1024**3}GB (建议至少8GB)"))
        
        # 磁盘空间检查 (至少50GB)
        if disk.free >= 50 * 1024**3:
            checks.append(("磁盘空间", "✅", f"{disk.free // 1024**3}GB 可用"))
        else:
            checks.append(("磁盘空间", "⚠️", f"{disk.free // 1024**3}GB (建议至少50GB)"))
        
        checks.append(("CPU核心", "✅", f"{cpu_count} 核心"))
        
        # 检查端口占用
        ports_to_check = [6379, 27017, 7474, 9200, 9092, 9000, 8000, 3000]
        occupied_ports = []
        
        for port in ports_to_check:
            if self._is_port_occupied(port):
                occupied_ports.append(port)
        
        if occupied_ports:
            checks.append(("端口", "⚠️", f"端口 {occupied_ports} 被占用"))
        else:
            checks.append(("端口", "✅", "所需端口可用"))
        
        # 打印检查结果表格
        table = Table(title="环境检查结果")
        table.add_column("项目", style="cyan")
        table.add_column("状态", style="magenta")
        table.add_column("详情", style="green")
        
        for check in checks:
            table.add_row(*check)
        
        self.console.print(table)
        
        return all(check[1] == "✅" for check in checks[:4])  # 前4项必须通过
    
    def _is_port_occupied(self, port: int) -> bool:
        """检查端口是否被占用"""
        import socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            result = sock.connect_ex(('localhost', port))
            return result == 0
        finally:
            sock.close()
    
    def setup_environment(self):
        """设置环境"""
        self.console.print("\n[bold blue]🔧 设置环境...[/bold blue]")
        
        # 创建必要目录
        directories = [
            'logs', 'data', 'models', 'static', 'temp',
            'data/raw', 'data/processed', 'data/models',
            'logs/crawler', 'logs/api', 'logs/ai', 'logs/celery'
        ]
        
        for directory in directories:
            os.makedirs(directory, exist_ok=True)
            self.logger.info(f"创建目录: {directory}")
        
        # 设置环境变量
        env_file = '.env'
        if not os.path.exists(env_file):
            self.console.print(f"[yellow]创建环境变量文件: {env_file}[/yellow]")
            with open(env_file, 'w') as f:
                f.write("""
# OpenAI API配置
OPENAI_API_KEY=your_openai_api_key_here
OPENAI_ORG_ID=your_openai_org_id_here

# Anthropic API配置
ANTHROPIC_API_KEY=your_anthropic_api_key_here

# GitHub API配置
GITHUB_TOKEN=your_github_token_here

# 数据库配置
MONGODB_URI=mongodb://admin:password123@localhost:27017/vulnerability_kg?authSource=admin
NEO4J_URI=bolt://localhost:7687
NEO4J_USERNAME=neo4j
NEO4J_PASSWORD=password123

# Redis配置
REDIS_URL=redis://localhost:6379/0

# Elasticsearch配置
ELASTICSEARCH_URL=http://localhost:9200

# Kafka配置
KAFKA_BOOTSTRAP_SERVERS=localhost:9092

# 安全配置
JWT_SECRET_KEY=your_jwt_secret_key_here
ENCRYPTION_KEY=your_encryption_key_here

# 监控配置
PROMETHEUS_URL=http://localhost:9090
GRAFANA_URL=http://localhost:3000

# 其他配置
DEBUG=true
LOG_LEVEL=INFO
MAX_WORKERS=4
BATCH_SIZE=100
""")
            self.console.print("[red]请编辑 .env 文件，填入正确的API密钥！[/red]")
    
    async def start_services(self, services: List[str] = None):
        """启动服务"""
        if services is None:
            services = ['infrastructure', 'applications']
        
        self.console.print("\n[bold blue]🚀 启动服务...[/bold blue]")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=self.console
        ) as progress:
            
            if 'infrastructure' in services:
                # 启动基础设施服务
                task = progress.add_task("启动基础设施服务...", total=None)
                
                result = subprocess.run([
                    'docker-compose', 'up', '-d',
                    'redis', 'mongodb', 'neo4j', 'elasticsearch', 
                    'zookeeper', 'kafka', 'minio'
                ], capture_output=True, text=True)
                
                if result.returncode != 0:
                    self.console.print(f"[red]基础设施启动失败: {result.stderr}[/red]")
                    return False
                
                progress.update(task, description="等待基础设施就绪...")
                await self._wait_for_services(self.services)
                progress.complete_task(task)
            
            if 'applications' in services:
                # 启动应用服务
                task = progress.add_task("启动应用服务...", total=None)
                
                result = subprocess.run([
                    'docker-compose', 'up', '-d',
                    'api-server', 'crawler-service', 'ai-processor',
                    'celery-worker', 'celery-beat'
                ], capture_output=True, text=True)
                
                if result.returncode != 0:
                    self.console.print(f"[red]应用服务启动失败: {result.stderr}[/red]")
                    return False
                
                progress.update(task, description="等待应用服务就绪...")
                await self._wait_for_services(self.app_services)
                progress.complete_task(task)
            
            if 'monitoring' in services:
                # 启动监控服务
                task = progress.add_task("启动监控服务...", total=None)
                
                result = subprocess.run([
                    'docker-compose', 'up', '-d',
                    'prometheus', 'grafana', 'nginx'
                ], capture_output=True, text=True)
                
                progress.complete_task(task)
        
        return True
    
    async def _wait_for_services(self, services: Dict):
        """等待服务就绪"""
        max_retries = 30
        retry_interval = 5
        
        for service_name, config in services.items():
            for attempt in range(max_retries):
                try:
                    if await config['healthcheck']():
                        self.logger.info(f"服务 {service_name} 就绪")
                        break
                except Exception as e:
                    if attempt == max_retries - 1:
                        self.logger.error(f"服务 {service_name} 启动失败: {e}")
                        raise
                    await asyncio.sleep(retry_interval)
    
    async def _check_redis(self) -> bool:
        """检查Redis服务"""
        import redis.asyncio as redis
        try:
            client = redis.Redis(host='localhost', port=6379, db=0)
            await client.ping()
            await client.close()
            return True
        except:
            return False
    
    async def _check_mongodb(self) -> bool:
        """检查MongoDB服务"""
        try:
            from motor.motor_asyncio import AsyncIOMotorClient
            client = AsyncIOMotorClient('mongodb://admin:password123@localhost:27017/')
            await client.admin.command('ping')
            client.close()
            return True
        except:
            return False
    
    async def _check_neo4j(self) -> bool:
        """检查Neo4j服务"""
        try:
            response = requests.get('http://localhost:7474/db/data/', timeout=5)
            return response.status_code == 200
        except:
            return False
    
    async def _check_elasticsearch(self) -> bool:
        """检查Elasticsearch服务"""
        try:
            response = requests.get('http://localhost:9200/_cluster/health', timeout=5)
            return response.status_code == 200
        except:
            return False
    
    async def _check_kafka(self) -> bool:
        """检查Kafka服务"""
        try:
            from kafka import KafkaProducer
            producer = KafkaProducer(
                bootstrap_servers=['localhost:9092'],
                value_serializer=lambda x: json.dumps(x).encode('utf-8')
            )
            producer.close()
            return True
        except:
            return False
    
    async def _check_minio(self) -> bool:
        """检查MinIO服务"""
        try:
            response = requests.get('http://localhost:9000/minio/health/live', timeout=5)
            return response.status_code == 200
        except:
            return False
    
    async def _check_api(self) -> bool:
        """检查API服务"""
        try:
            response = requests.get('http://localhost:8000/health', timeout=10)
            return response.status_code == 200
        except:
            return False
    
    async def _check_crawler(self) -> bool:
        """检查爬虫服务"""
        try:
            # 检查爬虫容器是否运行
            container = self.docker_client.containers.get('vulnerability-kg-crawler')
            return container.status == 'running'
        except:
            return False
    
    async def _check_ai_processor(self) -> bool:
        """检查AI处理服务"""
        try:
            container = self.docker_client.containers.get('vulnerability-kg-ai-processor')
            return container.status == 'running'
        except:
            return False
    
    async def _check_celery_worker(self) -> bool:
        """检查Celery Worker"""
        try:
            container = self.docker_client.containers.get('vulnerability-kg-celery-worker')
            return container.status == 'running'
        except:
            return False
    
    async def initialize_databases(self):
        """初始化数据库"""
        self.console.print("\n[bold blue]🗄️ 初始化数据库...[/bold blue]")
        
        tasks = [
            ("MongoDB", self._init_mongodb),
            ("Neo4j", self._init_neo4j),
            ("Elasticsearch", self._init_elasticsearch),
        ]
        
        for task_name, task_func in tasks:
            try:
                self.console.print(f"[yellow]初始化 {task_name}...[/yellow]")
                await task_func()
                self.console.print(f"[green]✅ {task_name} 初始化完成[/green]")
            except Exception as e:
                self.console.print(f"[red]❌ {task_name} 初始化失败: {e}[/red]")
    
    async def _init_mongodb(self):
        """初始化MongoDB"""
        from motor.motor_asyncio import AsyncIOMotorClient
        
        client = AsyncIOMotorClient('mongodb://admin:password123@localhost:27017/')
        db = client.vulnerability_kg
        
        # 创建集合和索引
        collections = [
            'vulnerabilities', 'products', 'attack_vectors',
            'mitigations', 'threat_actors', 'raw_data'
        ]
        
        for collection_name in collections:
            collection = db[collection_name]
            
            # 创建基本索引
            await collection.create_index("created_at")
            await collection.create_index("updated_at")
            
            # 创建特定索引
            if collection_name == 'vulnerabilities':
                await collection.create_index("cve_id", unique=True)
                await collection.create_index("severity")
                await collection.create_index("cvss_score")
            elif collection_name == 'products':
                await collection.create_index("cpe_id", unique=True)
                await collection.create_index("vendor")
            elif collection_name == 'raw_data':
                await collection.create_index("content_hash", unique=True)
                await collection.create_index("source")
        
        client.close()
    
    async def _init_neo4j(self):
        """初始化Neo4j"""
        from neo4j import AsyncGraphDatabase
        
        driver = AsyncGraphDatabase.driver(
            "bolt://localhost:7687",
            auth=("neo4j", "password123")
        )
        
        async with driver.session() as session:
            # 创建约束
            constraints = [
                "CREATE CONSTRAINT vuln_cve_id IF NOT EXISTS FOR (v:Vulnerability) REQUIRE v.cve_id IS UNIQUE",
                "CREATE CONSTRAINT product_cpe_id IF NOT EXISTS FOR (p:Product) REQUIRE p.cpe_id IS UNIQUE",
                "CREATE CONSTRAINT attack_technique_id IF NOT EXISTS FOR (a:AttackVector) REQUIRE a.technique_id IS UNIQUE"
            ]
            
            for constraint in constraints:
                await session.run(constraint)
            
            # 创建索引
            indexes = [
                "CREATE INDEX vuln_name_index IF NOT EXISTS FOR (v:Vulnerability) ON (v.name)",
                "CREATE INDEX product_name_index IF NOT EXISTS FOR (p:Product) ON (p.name)",
                "CREATE INDEX attack_name_index IF NOT EXISTS FOR (a:AttackVector) ON (a.name)"
            ]
            
            for index in indexes:
                await session.run(index)
        
        await driver.close()
    
    async def _init_elasticsearch(self):
        """初始化Elasticsearch"""
        import aiohttp
        
        # 创建索引模板
        index_template = {
            "index_patterns": ["vulnerability-*"],
            "template": {
                "mappings": {
                    "properties": {
                        "cve_id": {"type": "keyword"},
                        "title": {"type": "text", "analyzer": "standard"},
                        "description": {"type": "text", "analyzer": "standard"},
                        "severity": {"type": "keyword"},
                        "cvss_score": {"type": "float"},
                        "affected_products": {"type": "text"},
                        "created_at": {"type": "date"},
                        "content_vector": {"type": "dense_vector", "dims": 384}
                    }
                }
            }
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.put(
                'http://localhost:9200/_index_template/vulnerability-template',
                json=index_template
            ) as response:
                if response.status not in [200, 201]:
                    raise Exception(f"创建索引模板失败: {await response.text()}")
    
    def show_status(self):
        """显示系统状态"""
        self.console.print("\n[bold blue]📊 系统状态[/bold blue]")
        
        # 获取容器状态
        containers = self.docker_client.containers.list(all=True)
        vulnerability_containers = [
            c for c in containers 
            if c.name.startswith('vulnerability-kg-')
        ]
        
        table = Table(title="服务状态")
        table.add_column("服务", style="cyan")
        table.add_column("状态", style="magenta")
        table.add_column("端口", style="green")
        table.add_column("CPU", style="yellow")
        table.add_column("内存", style="blue")
        
        for container in vulnerability_containers:
            try:
                stats = container.stats(stream=False)
                cpu_percent = self._calculate_cpu_percent(stats)
                memory_usage = self._format_memory_usage(stats)
                
                status = "🟢 运行中" if container.status == 'running' else "🔴 已停止"
                ports = self._format_ports(container.ports)
                
                table.add_row(
                    container.name.replace('vulnerability-kg-', ''),
                    status,
                    ports,
                    f"{cpu_percent:.1f}%",
                    memory_usage
                )
            except Exception as e:
                table.add_row(
                    container.name.replace('vulnerability-kg-', ''),
                    "❌ 错误",
                    "-",
                    "-",
                    "-"
                )
        
        self.console.print(table)
        
        # 显示访问地址
        panel_content = """
🌐 访问地址:
• API文档: http://localhost:8000/docs
• Neo4j浏览器: http://localhost:7474
• Grafana监控: http://localhost:3000
• MinIO控制台: http://localhost:9001
• Elasticsearch: http://localhost:9200
• Web界面: http://localhost:3001

📚 管理命令:
• 查看日志: docker-compose logs -f [service]
• 重启服务: docker-compose restart [service]
• 停止所有: docker-compose down
• 清理数据: docker-compose down -v
        """
        
        self.console.print(Panel(panel_content, title="📖 使用指南", border_style="green"))
    
    def _calculate_cpu_percent(self, stats):
        """计算CPU使用率"""
        try:
            cpu_delta = stats['cpu_stats']['cpu_usage']['total_usage'] - \
                       stats['precpu_stats']['cpu_usage']['total_usage']
            system_delta = stats['cpu_stats']['system_cpu_usage'] - \
                          stats['precpu_stats']['system_cpu_usage']
            
            if system_delta > 0:
                return (cpu_delta / system_delta) * len(stats['cpu_stats']['cpu_usage']['percpu_usage']) * 100.0
        except:
            pass
        return 0.0
    
    def _format_memory_usage(self, stats):
        """格式化内存使用"""
        try:
            usage = stats['memory_stats']['usage']
            limit = stats['memory_stats']['limit']
            return f"{usage // 1024 // 1024}MB / {limit // 1024 // 1024}MB"
        except:
            return "-"
    
    def _format_ports(self, ports):
        """格式化端口信息"""
        if not ports:
            return "-"
        
        port_list = []
        for internal_port, external_ports in ports.items():
            if external_ports:
                external_port = external_ports[0]['HostPort']
                port_list.append(f"{external_port}:{internal_port}")
        
        return ", ".join(port_list) if port_list else "-"
    
    async def stop_services(self):
        """停止服务"""
        self.console.print("\n[bold red]🛑 停止服务...[/bold red]")
        
        result = subprocess.run(['docker-compose', 'down'], 
                              capture_output=True, text=True)
        
        if result.returncode == 0:
            self.console.print("[green]✅ 所有服务已停止[/green]")
        else:
            self.console.print(f"[red]❌ 停止服务失败: {result.stderr}[/red]")
    
    async def cleanup(self):
        """清理系统"""
        self.console.print("\n[bold red]🧹 清理系统...[/bold red]")
        
        # 停止并删除容器，清理数据卷
        result = subprocess.run(['docker-compose', 'down', '-v', '--remove-orphans'], 
                              capture_output=True, text=True)
        
        if result.returncode == 0:
            self.console.print("[green]✅ 系统清理完成[/green]")
        else:
            self.console.print(f"[red]❌ 清理失败: {result.stderr}[/red]")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='网络安全漏洞知识图谱项目管理器')
    parser.add_argument('command', choices=[
        'check', 'setup', 'start', 'stop', 'status', 'cleanup', 'init-db'
    ], help='要执行的命令')
    parser.add_argument('--services', nargs='+', help='指定要启动的服务组')
    
    args = parser.parse_args()
    
    manager = ProjectManager()
    
    try:
        if args.command == 'check':
            success = manager.check_prerequisites()
            sys.exit(0 if success else 1)
            
        elif args.command == 'setup':
            manager.setup_environment()
            
        elif args.command == 'start':
            if not manager.check_prerequisites():
                manager.console.print("[red]❌ 环境检查失败，请先解决问题[/red]")
                sys.exit(1)
            
            manager.setup_environment()
            success = await manager.start_services(args.services)
            
            if success:
                await manager.initialize_databases()
                manager.show_status()
            else:
                sys.exit(1)
                
        elif args.command == 'stop':
            await manager.stop_services()
            
        elif args.command == 'status':
            manager.show_status()
            
        elif args.command == 'cleanup':
            await manager.cleanup()
            
        elif args.command == 'init-db':
            await manager.initialize_databases()
            
    except KeyboardInterrupt:
        manager.console.print("\n[yellow]⚠️ 操作被用户中断[/yellow]")
    except Exception as e:
        manager.console.print(f"[red]❌ 执行失败: {e}[/red]")
        sys.exit(1)


if __name__ == '__main__':
    asyncio.run(main()) 